import numpy as np

# 实现树的结点类
class Node:
    def __init__(self, data, lchild=None, rchild=None):
        self.data = data
        self.lchild = lchild
        self.rchild = rchild


class kdTree:
    def __init__(self):
        self.kdTree = None

    # 创建kd树
    def create(self, dataSet, depth):
        if (len(dataSet) > 0):
            m, n = np.shape(dataSet)
            midIndex = int(m / 2) # 中间树的索引位置
            axis = depth % n # 判断以哪个轴划分数据
            sortedDataSet = self.sort(dataSet, axis) # 对数据进行排序
            node = Node(sortedDataSet[midIndex]) # 将结点数据域设置为中位数
            # print("中位数：")
            # print(sortedDataSet[midIndex])
            leftDataSet = sortedDataSet[: midIndex]
            rightDataSet = sortedDataSet[midIndex + 1 :] # 将大于中位数的归入右子树，小于的归入左子树
            print("左子树：")
            print(leftDataSet)
            print("右子树")
            print(rightDataSet)
            node.lchild = self.create(leftDataSet, depth+1) # 将中位数左边样本传入来递归创建树
            node.rchild = self.create(rightDataSet, depth+1) # 同上
            return node
        else:
            return None


    # 冒泡排序
    def sort(self, dataSet, axis):
        sortDataSet = dataSet[:] # 为了不破坏原样本，建立此副本
        m, n = np.shape(sortDataSet)
        for i in range(m):
            for j in range(0, m - i - 1):
                if (sortDataSet[j][axis] > sortDataSet[j + 1][axis]):
                    t = sortDataSet[j]
                    sortDataSet[j] = sortDataSet[j + 1]
                    sortDataSet[j + 1] = t
        print("sortDataSet:")
        print(sortDataSet)
        return sortDataSet


    # 先序遍历检查树有无问题
    def preOrder(self, node):
        if node != None:
            print("preOrder-->%s" % node.data)
            self.preOrder(node.lchild)
            self.preOrder(node.rchild)


    # kdTree的搜索
    def search(self, tree, x):
        self.nearestPoint = None # 保存最近的点
        self.nearestValue = 0 # 保存最近点的值（距离）
        def travel(node, depth=0): # 递归搜索
            if node != None:  # 递归终止条件
                n = len(x) # 特征数
                axis = depth % n # 计算用于分割的轴
                print("*" * 50)
                print(node.data[axis])
                print("*" * 50)
                if x[axis] < node.data[axis]: # 如果数据小于结点，则在左节点找
                    travel(node.lchild, depth+1)
                else:
                    travel(node.rchild, depth+1)

                # 递归结束，向上回退
                distNodeAndX = self.dist(x, node.data) # 目标和结点的距离判断
                if (self.nearestPoint == None): # 确定当前最近点
                    self.nearestPoint = node.data
                    self.nearestValue = distNodeAndX
                elif (self.nearestValue > distNodeAndX):
                    self.nearestPoint = node.data
                    self.nearestValue = distNodeAndX

                print(node.data, depth, self.nearestValue, node.data[axis], x[axis])
                # 确定是否需要去子节点的区域去找（圆的判断）
                if (abs(x[axis] - node.data[axis]) <= self.nearestValue):
                    if (x[axis] < node.data[axis]):
                        travel(node.rchild, depth+1)
                    else:
                        travel(node.lchild, depth+1)
        travel(tree)
        return self.nearestPoint

    def dist(self, x1, x2):
        return ((np.array(x1) - np.array(x2)) ** 2).sum() ** 0.5


dataSet = [[2, 3],
           [5, 4],
           [9, 6],
           [4, 7],
           [8, 1],
           [7, 2]]
x = [5, 3]
kdtree = kdTree()
tree = kdtree.create(dataSet, 0)
kdtree.preOrder(tree)
print(kdtree.search(tree, x))


